package cn.tsingmo.otproxy.security;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelOutboundHandlerAdapter;
import io.netty.channel.ChannelPromise;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 加密处理器
 * 负责在数据离开管道时对其进行加密
 */
public class EncryptionHandler extends ChannelOutboundHandlerAdapter {
    private static final Logger logger = LoggerFactory.getLogger(EncryptionHandler.class);
    private final EncryptionService encryptionService;

    /**
     * 构造函数
     * @param encryptionService 加密服务实例
     */
    public EncryptionHandler(EncryptionService encryptionService) {
        this.encryptionService = encryptionService;
    }

    /**
     * 写入数据并加密
     */
    @Override
    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
        if (!(msg instanceof byte[])) {
            logger.error("发送非字节数组类型的消息");
            ctx.close();
            return;
        }

        try {
            // 加密数据
            byte[] data = (byte[]) msg;
            byte[] encryptedData = encryptionService.encrypt(data);

            // 将加密后的数据包装为ByteBuf并发送
            ByteBuf byteBuf = Unpooled.wrappedBuffer(encryptedData);
            ctx.writeAndFlush(byteBuf, promise);
        } catch (Exception e) {
            logger.error("数据加密失败", e);
            ctx.close(); // 加密失败时关闭连接
            promise.setFailure(e);
        }
    }
}